#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <xcb/xcb.h>
#include <xcb/dri3.h>
#include <xcb/present.h>
#include <xcb/xcb_icccm.h>
#include <sys/mman.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <time.h>
#include <errno.h>
#include <stdint.h>
#include <assert.h>
#include <drm/drm_fourcc.h>
#include <sys/ioctl.h>
#include <xf86drm.h>
#include <xf86drmMode.h>
#include <drm/drm.h>
#include <drm/drm_mode.h>

#define WIDTH  400
#define HEIGHT 300

static int create_shm_file(size_t size) {
    char name[] = "/tmp/x11_dri3_shm-XXXXXX";
    int fd = mkstemp(name);
    if (fd < 0) return -1;
    unlink(name);
    if (ftruncate(fd, size) < 0) {
        close(fd);
        return -1;
    }
    return fd;
}

static void fill_color(uint32_t *buf, int width, int height, uint32_t color) {
    for (int i = 0; i < width * height; ++i)
        buf[i] = color;
}

int main() {
    xcb_connection_t *conn = xcb_connect(NULL, NULL);
    if (xcb_connection_has_error(conn)) {
        fprintf(stderr, "Failed to connect to X server\n");
        return 1;
    }

    const xcb_setup_t *setup = xcb_get_setup(conn);
    xcb_screen_t *screen = xcb_setup_roots_iterator(setup).data;

    xcb_window_t win = xcb_generate_id(conn);
    uint32_t mask = XCB_CW_BACK_PIXEL | XCB_CW_EVENT_MASK;
    uint32_t values[] = { screen->black_pixel, XCB_EVENT_MASK_EXPOSURE };

    xcb_create_window(conn, XCB_COPY_FROM_PARENT, win, screen->root,
                      0, 0, WIDTH, HEIGHT, 0,
                      XCB_WINDOW_CLASS_INPUT_OUTPUT, screen->root_visual,
                      mask, values);

    xcb_icccm_set_wm_name(conn, win, XCB_ATOM_STRING, 8, strlen("X11 DRI3 Color Switch"), "X11 DRI3 Color Switch");
    xcb_map_window(conn, win);
    xcb_flush(conn);

    // Query DRI3 support
    xcb_dri3_open_cookie_t open_cookie = xcb_dri3_open(conn, win,  0);
    xcb_dri3_open_reply_t *open_reply = xcb_dri3_open_reply(conn, open_cookie, NULL);
    if (!open_reply || open_reply->nfd == 0) {
        fprintf(stderr, "DRI3 not available or failed to open\n");
        free(open_reply);
        xcb_disconnect(conn);
        return 1;
    }
    int drm_fd = xcb_dri3_open_reply_fds(conn, open_reply)[0];
    free(open_reply);

    // Create a dumb buffer
    struct drm_mode_create_dumb creq = {0};
    creq.width = WIDTH;
    creq.height = HEIGHT;
    creq.bpp = 32;
    if (ioctl(drm_fd, DRM_IOCTL_MODE_CREATE_DUMB, &creq) < 0) {
        perror("DRM_IOCTL_MODE_CREATE_DUMB");
        close(drm_fd);
        xcb_disconnect(conn);
        return 1;
    }

    // Export the buffer as a PRIME fd
    struct drm_prime_handle prime = {0};
    prime.handle = creq.handle;
    prime.flags = DRM_CLOEXEC | DRM_RDWR;
    if (ioctl(drm_fd, DRM_IOCTL_PRIME_HANDLE_TO_FD, &prime) < 0) {
        perror("DRM_IOCTL_PRIME_HANDLE_TO_FD");
        close(drm_fd);
        xcb_disconnect(conn);
        return 1;
    }
    int prime_fd = prime.fd;

    // Map the buffer
    struct drm_mode_map_dumb mreq = {0};
    mreq.handle = creq.handle;
    if (ioctl(drm_fd, DRM_IOCTL_MODE_MAP_DUMB, &mreq) < 0) {
        perror("DRM_IOCTL_MODE_MAP_DUMB");
        close(prime_fd);
        close(drm_fd);
        xcb_disconnect(conn);
        return 1;
    }
    uint32_t *buf = mmap(0, creq.size, PROT_READ | PROT_WRITE, MAP_SHARED, drm_fd, mreq.offset);
    if (buf == MAP_FAILED) {
        perror("mmap");
        close(prime_fd);
        close(drm_fd);
        xcb_disconnect(conn);
        return 1;
    }

    // Create pixmap from PRIME fd
    xcb_pixmap_t pixmap = xcb_generate_id(conn);
    xcb_dri3_pixmap_from_buffer(conn, pixmap, win, creq.size, WIDTH, HEIGHT, creq.pitch, 32, 24, prime_fd);

    // Present extension
    xcb_present_select_input(conn, xcb_generate_id(conn), win, XCB_PRESENT_EVENT_MASK_COMPLETE_NOTIFY);

    int toggle = 0;
    uint32_t colors[2] = { 0xff0000ff, 0xffff0000 }; // blue, red

    while (1) {
        fill_color(buf, WIDTH, HEIGHT, colors[toggle]);
        xcb_copy_area(conn, pixmap, win, screen->black_pixel, 0, 0, 0, 0, WIDTH, HEIGHT);
        xcb_flush(conn);
        toggle = !toggle;

        // Wait for 1 second or X event
        fd_set fds;
        FD_ZERO(&fds);
        FD_SET(xcb_get_file_descriptor(conn), &fds);
        struct timeval tv = {1, 0};
        int ret = select(xcb_get_file_descriptor(conn)+1, &fds, NULL, NULL, &tv);
        if (ret > 0) {
            xcb_generic_event_t *ev;
            while ((ev = xcb_poll_for_event(conn))) {
                if ((ev->response_type & ~0x80) == XCB_EXPOSE) {
                    // Redraw on expose
                    xcb_copy_area(conn, pixmap, win, screen->black_pixel, 0, 0, 0, 0, WIDTH, HEIGHT);
                    xcb_flush(conn);
                }
                free(ev);
            }
        }
    }

    munmap(buf, creq.size);
    close(prime_fd);
    close(drm_fd);
    xcb_disconnect(conn);
    return 0;
}
